pub mod config;

use network_scanner_server_ipc::RecordServiceClient;
use rustls::pki_types::ServerName;
use rustls::RootCertStore;
use std::net::ToSocketAddrs;
use std::sync::Arc;
use tarpc::client::Config;
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;

use rustls_pemfile;
use tarpc::tokio_serde::formats::Bincode;

use self::config::RpcConfig;

// 定义RPC连接函数，用于实现远程连接
pub async fn connect(config: RpcConfig) -> RecordServiceClient {
    let mut root_cert_store = RootCertStore::empty();
    let mut cert_chain = Vec::new();

    // 读取客户端证书
    let client_cert = rustls_pemfile::read_one_from_slice(config.client_certificate.as_bytes())
        .expect("Read client certificate failed")
        .expect("No client certificate provided");
    let client_cert = match client_cert.0 {
        rustls_pemfile::Item::X509Certificate(cert) => cert,
        _ => panic!("Unexpected item in client certificate"),
    };
    cert_chain.push(client_cert);

    // 读取服务端证书
    let mut server_cert = config.server_certificate.as_bytes();
    loop {
        let cert = rustls_pemfile::read_one_from_slice(server_cert)
            .expect("Read server certificate failed");
        let cert = match cert {
            Some(cert) => cert,
            None => break,
        };
        server_cert = cert.1;
        let cert = match cert.0 {
            rustls_pemfile::Item::X509Certificate(cert) => cert,
            _ => panic!("unexpected item in cert chain"),
        };
        if root_cert_store.is_empty() {
            root_cert_store
                .add(cert.clone())
                .expect("Failed to add certificate into root certificate store");
        }
        cert_chain.push(cert);
    }

    if root_cert_store.is_empty() {
        panic!("No root ca in config")
    }

    // 读取客户端私钥
    let client_key = rustls_pemfile::read_one_from_slice(config.client_key.as_bytes())
        .expect("Read client key failed")
        .expect("No client key provided");
    let client_key = match client_key.0 {
        rustls_pemfile::Item::Pkcs8Key(client_key) => client_key,
        _ => panic!("Unexpected item in client key"),
    };

    // 创建TLS连接
    let client_config = rustls::ClientConfig::builder()
        .with_root_certificates(root_cert_store)
        .with_client_auth_cert(cert_chain, client_key.into())
        .expect("Client key is invalid");
    let connector = TlsConnector::from(Arc::new(client_config));

    let addr = config
        .host
        .to_socket_addrs()
        .expect("Unable to resolve server host")
        .next()
        .expect("No address resolved for host");

    let stream = TcpStream::connect(addr).await.unwrap();

    // 创建RPC连接
    let mut transport = tarpc_rustls::tls::connect(
        connector,
        stream,
        ServerName::DnsName("network-scanner-server".try_into().unwrap()),
        Bincode::default,
    );
    transport.config_mut().max_frame_length(1024usize);

    // 返回RPC客户端
    RecordServiceClient::new(Config::default(), transport.await.unwrap()).spawn()
}
